# -*- coding:utf-8 -*-

# @Time    : 18-11-17 下午12:16

# @Author  : Swing


from tensorflow.examples.tutorials.mnist import input_data

import tensorflow as tf

# Import data
# data_dir = '/home/swing/Documents/data/mnist/'
# data_dir = 'data/'
data_dir = '/Users/zhubin/Documents/ai/data/mnist/'
mnist = input_data.read_data_sets(data_dir, one_hot=True)
# Define loss and optimizer
x = tf.placeholder(tf.float32, [None, 784], name='x')
y_ = tf.placeholder(tf.float32, [None, 10], name='y_')
learning_rate = tf.placeholder(tf.float32)

with tf.name_scope('reshape'):
    x_image = tf.reshape(x, [-1, 28, 28, 1], name='x_image')

# First convolutional layer - maps one grayscale image to 32 feature maps.
with tf.name_scope('conv1'):
    shape = [5, 5, 1, 32]
    w_conv1 = tf.Variable(tf.truncated_normal(shape, stddev=0.1, name='w_conv1'),
                          collections=[tf.GraphKeys.GLOBAL_VARIABLES, 'WEIGHTS'])
    shape = [32]
    b_conv1 = tf.Variable(tf.constant(0.1, shape=shape), name='b_conv1')
    l_conv1 = tf.nn.conv2d(x_image, w_conv1, strides=[1, 1, 1, 1], padding='SAME', name='l_conv1') + b_conv1
    h_conv1 = tf.nn.relu(l_conv1, name='h_conv1')

# Pooling layer - downsamples by 2X.
with tf.name_scope('pool1'):
    h_pool1 = tf.nn.max_pool(h_conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID', name='h_pool1')


with tf.name_scope('conv2'):
    w_conv2 = tf.Variable(tf.truncated_normal([5, 5, 32, 64], stddev=0.1),
                          collections=[tf.GraphKeys.GLOBAL_VARIABLES, 'WEIGHTS'],
                          name='w_conv2')

    b_conv2 = tf.Variable(tf.constant(0.1, shape=[64]), name='b_conv2')

    l_conv2 = tf.nn.conv2d(h_pool1, w_conv2, strides=[1, 1, 1, 1], padding='SAME') + b_conv2

    h_conv2 = tf.nn.relu(l_conv2)

with tf.name_scope('pool2'):
    h_pool2 = tf.nn.max_pool(h_conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID', name='h_pool2')

with tf.name_scope('fc1'):
    w_fc1 = tf.Variable(tf.truncated_normal([7 * 7 * 64, 1024], stddev=0.1, name='w_fc1'),
                        collections=[tf.GraphKeys.GLOBAL_VARIABLES, 'WEIGHTS'])
    b_fc1 = tf.Variable(tf.constant(0.1, shape=[1024]), name='b_fc1')
    h_pool2_flat = tf.reshape(h_pool2, shape=[-1, 7 * 7 * 64])
    h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, w_fc1) + b_fc1)

with tf.name_scope('dropout'):
    keepProb = tf.placeholder(tf.float32)
    h_fc1_drop = tf.nn.dropout(x=h_fc1, keep_prob=keepProb)

with tf.name_scope('fc2'):
    w_fc2 = tf.Variable(tf.truncated_normal([1024, 10], stddev=0.1, name='w_fc2'),
                        collections=[tf.GraphKeys.GLOBAL_VARIABLES, 'WEIGHTS'], name='b_fc2')
    b_fc2 = tf.Variable(tf.constant(0.1, shape=[10]))

    y = tf.matmul(h_fc1_drop, w_fc2) + b_fc2


cross_entropy = tf.reduce_mean(
    tf.nn.softmax_cross_entropy_with_logits(labels=y_, logits=y)
)

l2_loss = tf.add_n([tf.nn.l2_loss(w) for w in tf.get_collection('WEIGHTS')])
total_loss = cross_entropy + 7e-5 * l2_loss

train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(total_loss)

sess = tf.Session()
init_op = tf.global_variables_initializer()
sess.run(init_op)

correct_predict = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_predict, tf.float32))
lr = 0.01
# train
for step in range(3000):
    batch_xs, batch_ys = mnist.train.next_batch(100)

    _, learning_rate_value, cross_entropy_value, l2_loss_value, total_loss_value = sess.run([train_step, learning_rate, cross_entropy, l2_loss, total_loss],
                                                                                            feed_dict={
                                                                                                x: batch_xs,
                                                                                                y_: batch_ys,
                                                                                                learning_rate: lr,
                                                                                                keepProb: 0.5
                                                                                            })

    if (step + 1) % 100 == 0:
        print('step %d, learniong rate %f entropy loss %f, l2 loss %f, total loss %f.'
              % (step + 1, learning_rate_value, cross_entropy_value, l2_loss_value, total_loss_value))
        print('accuracy: ', sess.run(accuracy, feed_dict={
            x: batch_xs,
            y_: batch_ys,
            keepProb: 0.5
        }))

    if (step + 1) % 1000 == 0:
        print('Test accuracy: ', sess.run(accuracy, feed_dict={
            x: mnist.test.images,
            y_: mnist.test.labels,
            keepProb: 0.5
        }))
